import numpy as np
import matplotlib.pyplot as plt
from disba import PhaseDispersion
import bayesbay as bb
from bayesbay.discretization import Voronoi1D
import os
from functools import partial

# 创建结果保存目录
output_dir = "./"
os.makedirs(output_dir, exist_ok=True)

# 读取频散曲线数据
def read_dispersion_curves(filepath):
    """
    从文件中读取频散曲线数据
    """
    with open(filepath, 'r') as f:
        lines = f.readlines()
    
    curves = []
    current_curve = None
    
    for line in lines:
        line = line.strip()
        if line.startswith('#'):
            if 'Curve' in line:
                if current_curve is not None:
                    curves.append(np.array(current_curve))
                current_curve = []
        else:
            if line and current_curve is not None:
                freq, vel = map(float, line.split())
                current_curve.append([freq, vel])
    
    # 添加最后一个curve
    if current_curve is not None:
        curves.append(np.array(current_curve))
    
    return curves

# 读取数据
curves = read_dispersion_curves('/home/shiyongxiang/coda_hsr/script.inv/P264S.D123.Tp.N.F1.0.3.5.V100.1400.NORM.STACK.FOLDER.disp1')

# 分别提取两个模式的数据
mode0_data = curves[0]  # 第一阶频散曲线
mode1_data = curves[1]  # 第二阶频散曲线

# 提取频率和相速度
freq0, d_r0 = mode0_data[-1::-1, 0], mode0_data[-1::-1, 1]/1000
freq1, d_r1 = mode1_data[-1::-1, 0], mode1_data[-1::-1, 1]/1000

# 将频率转换为周期
periods0 = 1.0 / freq0
periods1 = 1.0 / freq1

# 定义参数化模型
def initialize_vs(param, positions=None):
    vmin, vmax = param.get_vmin_vmax(positions)
    sorted_vals = np.sort(np.random.uniform(vmin, vmax, positions.size))
    return sorted_vals

vs = bb.prior.UniformPrior(name="vs", 
                           vmin=[0.2, 0.3, 0.4], 
                           vmax=[0.5, 0.6, 1.4], 
                           position=[0, 0.1, 0.2],
                           perturb_std=0.03)

vs.set_custom_initialize(initialize_vs)

vp_vs = 1.77  # vp/vs比率
rho_vp = [0.32, 0.77]  # rho = 0.32 * vp + 0.77
Zmin,Zmax = 0,0.4
voronoi = Voronoi1D(
    name="voronoi", 
    vmin=Zmin,   #最小深度
    vmax=Zmax, #最大深度
    perturb_std=0.01,
    n_dimensions=None, 
    n_dimensions_min=3,
    n_dimensions_max=10,
    parameters=[vs], 
    birth_from='neighbour'
)

parameterization = bb.parameterization.Parameterization(voronoi)

# 正演函数
def forward_sw(state, periods, wave="rayleigh", mode=0):
    voronoi = state["voronoi"]
    voronoi_sites = voronoi["discretization"]
    thickness = Voronoi1D.compute_cell_extents(voronoi_sites)
    vs = voronoi["vs"]
    vp = vs * vp_vs
    rho = rho_vp[0] * vp + rho_vp[1]
    
    # 计算瑞雷波相速度
    pd = PhaseDispersion(thickness, vp, vs, rho)
    d_pred_rayleigh = pd(periods, mode=mode, wave="rayleigh").velocity

    if len(d_pred_rayleigh) < len(periods):
        re = np.zeros_like(len(periods)-len(d_pred_rayleigh))+max(vs)
        d_pred_rayleigh = np.concatenate([d_pred_rayleigh, re])

    return d_pred_rayleigh

target_rayleigh = bb.Target("R0", 
                            d_r0, 
                            std_min=0.001, 
                            std_max=0.01, 
                            std_perturb_std=0.001)
target_rayleigh1 = bb.Target("R1", 
                            d_r1, 
                            std_min=0.001, 
                            std_max=0.01, 
                            std_perturb_std=0.001)
f1 = partial(forward_sw, periods = periods0, wave="rayleigh", mode=0)
f2 = partial(forward_sw, periods = periods1, wave="rayleigh", mode=1)
log_likelihood = bb.LogLikelihood(targets=[target_rayleigh,target_rayleigh1], fwd_functions=[f1,f2])
# log_likelihood = bb.LogLikelihood(targets=[target_rayleigh1], fwd_functions=[forward_sw1])

# 贝叶斯反演
inversion = bb.BayesianInversion(
    parameterization=parameterization, 
    log_likelihood=log_likelihood,
    n_chains=32
)

print('观测数据参数')
print(periods0.shape)
print(d_r0.shape)
print(periods1.shape)
print(d_r1.shape)

print("开始反演...")
inversion.run(
    sampler=None, 
    n_iterations=200_000, 
    burnin_iterations=20_000, 
    save_every=100,
    verbose=False,
    print_every=4000,
)

# 输出统计信息
for chain in inversion.chains:
    chain.print_statistics()

# 获取结果
results = inversion.get_results(concatenate_chains=True)


# 处理结果
sampled_voronoi_nuclei = results['voronoi.discretization']
sampled_thickness = [Voronoi1D.compute_cell_extents(n) for n in sampled_voronoi_nuclei]
sampled_vs = results['voronoi.vs']
interp_depths = np.linspace(Zmin, Zmax, 100)
statistics_vs = Voronoi1D.get_tessellation_statistics(
    sampled_thickness, sampled_vs, interp_depths, input_type='extents'
)

sampled_thickness_valid = []
for i in sampled_thickness:
    ti = list(i[:-1])+[0.05]
    sampled_thickness_valid.append(np.array(ti))
sampled_thickness = sampled_thickness_valid


# 绘制结果
fig, axes = plt.subplots(1, 2, figsize=(12, 6))

# 绘制频散曲线对比
axes[0].plot(1/periods0, d_r0, 'ro', label='R0')
d_pred_rayleigh_all = np.array(results['R0.dpred'])
percentiles_rayleigh = np.percentile(d_pred_rayleigh_all, (16, 84), axis=0)
axes[0].fill_between(1/periods0, percentiles_rayleigh[0], percentiles_rayleigh[1], 
                     color='red', alpha=0.3, label='Predicted Rayleigh (mean ± 1σ)')

axes[0].plot(1/periods1, d_r1, 'ro', label='R1')
d_pred_rayleigh_all = np.array(results['R1.dpred'])
percentiles_rayleigh = np.percentile(d_pred_rayleigh_all, (16, 84), axis=0)
axes[0].fill_between(1/periods1, percentiles_rayleigh[0], percentiles_rayleigh[1], 
                     color='red', alpha=0.3, label='Predicted Rayleigh (mean ± 1σ)')

axes[0].set_xlabel('F(Hz)')
axes[0].set_ylabel('Phase velocity [km/s]')
axes[0].legend()
axes[0].grid(True)
axes[0].set_title('Dispersion Curves')

# 绘制Vs结构
ax1 = axes[1]
Voronoi1D.plot_tessellation_density(sampled_thickness, sampled_vs, 
                                    input_type='extents', ax=ax1, cmap='gray_r')
ax1.plot(statistics_vs['mean'], interp_depths, 'r-', lw=2, label='Mean')
ax1.plot(statistics_vs['median'], interp_depths, 'b--', lw=2, label='Median')
ax1.set_xlabel('Vs [km/s]')
ax1.set_ylabel('Depth [km]')
ax1.legend()
ax1.grid(True)
ax1.set_title('Inverted Vs Profile')

plt.tight_layout()
plt.savefig(os.path.join(output_dir, "inversion_results.png"), dpi=300, bbox_inches='tight')

def get_subplot_layout(n_subplots):
    rows = int(np.sqrt(n_subplots))
    cols = int(np.ceil(n_subplots / rows))
    return rows, cols

rows, cols = get_subplot_layout(len(inversion.chains))
fig, axes = plt.subplots(rows, cols, figsize=(15, 15))
for ipanel, (ax, chain) in enumerate(zip(np.ravel(axes), inversion.chains)):
    saved_states = chain.saved_states
    saved_nuclei = saved_states["voronoi.discretization"]
    saved_vs = saved_states['voronoi.vs']
    
    Voronoi1D.plot_tessellations(saved_nuclei,
                                 saved_vs, 
                                 ax=ax, 
                                 linewidth=0.1, 
                                 color="k", 
                                 bounds=(0, Zmax))
    Voronoi1D.plot_tessellation_statistics(
        saved_nuclei, saved_vs, interp_depths, ax=ax
    )
    
    ax.set_title(f'Chain {chain.id}')
    ax.tick_params(direction='in', labelleft=False, labelbottom=False)
    ax.set_xlabel('')
    ax.set_ylabel('')
    
    if not ipanel % cols:
        ax.set_ylabel('Depth [km]')
        ax.tick_params(labelleft=True)
    if ipanel >= (rows-1) * cols:
        ax.set_xlabel('Vs [km/s]')
        ax.tick_params(labelbottom=True)
    
plt.tight_layout()
plt.savefig(os.path.join(output_dir, "inversion_chains.png"), dpi=300, bbox_inches='tight')

